import os
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

from pdf_processor import PDFProcessor
from vector_store import VectorStore

class PDFQASystem:
    def __init__(self):
        self.vector_store = VectorStore()
        self.confidence_threshold = 60  # OCR置信度阈值

    def process_pdf(self, pdf_path: str):
        """处理PDF文件并存储到向量数据库"""
        processor = PDFProcessor(pdf_path)
        document_data = processor.process_document()
        
        # 处理文本块
        text_documents = [
            {
                'content': chunk,
                'type': 'text',
                'page_num': i // 1000 + 1,  # 估算页码
                'confidence': 100  # 文本内容置信度为100%
            }
            for i, chunk in enumerate(document_data['text_chunks'])
        ]
        
        # 处理图片
        image_documents = [
            {
                'content': f"图片内容：{img['text']}",
                'type': 'image',
                'page_num': img['page'],
                'confidence': img['confidence']
            }
            for img in document_data['images']
        ]
        
        # 合并所有文档并存储
        all_documents = text_documents + image_documents
        self.vector_store.add_documents(all_documents)

    def answer_question(self, question: str) -> str:
        """回答用户问题"""
        # 搜索相关文档
        relevant_docs = self.vector_store.search(question)
        
        # 根据置信度过滤和排序文档
        filtered_docs = [
            doc for doc in relevant_docs 
            if doc.get('confidence', 100) >= self.confidence_threshold
        ]
        
        if not filtered_docs:
            return "抱歉，没有找到足够可信的相关信息来回答您的问题。"
        
        # 构建上下文，优先使用高置信度的内容
        context = "\n".join([
            f"{doc['content']} (置信度: {doc.get('confidence', 100)}%)"
            for doc in sorted(filtered_docs, key=lambda x: x.get('confidence', 100), reverse=True)
        ])

        llm = ChatOpenAI(
            api_key="sk-a3f7718fb81f43b2915f0a6483b6661b",
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
            model="llama-4-scout-17b-16e-instruct",
            # 此处以qwen-plus为例，您可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
            # other params...
        )
        template = """
        你是一个专业的问题助手，结合上下文简要的回答问题，但是在答案里面请不要提及根据提供的信息等关键词，只需要直接回答出问题的答案即可，如果未匹配到正确答案请回答：我无法回答您的问题。
        上下文:{content}
        问题:{question}
        """
        str = StrOutputParser()
        prompt = ChatPromptTemplate.from_template(template)
        chain = prompt | llm | str
        res = chain.invoke({"content": context, "question": question})
        print(res)
        return res

def main():
    # 示例使用
    qa_system = PDFQASystem()
    
    # 处理PDF文件
    pdf_path = r"D:\test_document.pdf"  # 替换为实际的PDF文件路径
    if os.path.exists(pdf_path):
        qa_system.process_pdf(pdf_path)
        print("PDF处理完成！")
    
    # 交互式问答
    while True:
        question = input("\n请输入您的问题（输入'退出'结束）：")
        if question.lower() in ['退出', 'exit', 'quit']:
            break
        
        answer = qa_system.answer_question(question)
        print("\n" + answer)

if __name__ == "__main__":
    main() 